import astro.fit
import astro.spec
from astro.utilities import Bunch
import astro.wrappers
import astro.coord
import os
import matplotlib.pyplot as pl
import numpy as np
#from astro.utilities import ismember
#import gas_gal_hiz.read
import astro.aaomega as aaom
#from copy import deepcopy
from io import readtxt


def readqso_info(filename):
    print 'Reading QSO information'
    fh = open(filename)
    names = (i.strip().replace('_j2000deg','') for i in fh.next().split('|'))
    names = [n.replace('magR','mag') for n in names]
    fh.close()
    qsos = readtxt(filename, sep='|', skip=1, names=names)
    #cqsos = qsos[qsos.z > 2.25]
    return qsos

def match_spectra_to_qsos(aaom_spec, qsos, arcsec=2):
    """ Find an aaomega spectra for each qso.

    Return a tuple of all te spectra (and their info) for each qso."""
    spectra = [[] for i in range(len(qsos))]
    for spec,info in aaom_spec:
        ra,dec = info.RA, info.DEC
        imatch = astro.coord.match(qsos.ra, qsos.dec, ra, dec, arcsec)
        for i,q in enumerate(qsos):
            ind = imatch[i]
            if ind != -1:
                spectra[i].append((spec[ind], info[ind]))
    for i,sp in enumerate(spectra):
        if sp:
            spec = [s[0] for s in sp]
            info = [s[1] for s in sp]
            spectra[i] = (spec, info)
            #spectra[i] = zip(*sp)
        else:
            spectra[i] = ([],[])

    return spectra

def getgrat(wa):
    """ Given a wavelength scale, determine which AAOmega grating it
    corresponds to.

    Heuristic.
    """
    if wa.ptp() > 1200:
        if wa.max() < 6200:
            return '580V' 
        else:
            return '385R'
    elif wa.max() > 6200:
        return '1000R'
    elif wa.max() > 5200:
        return '1500Vr'
    else:
        return '1500Vb'

def scaling(spec, bscale, rscale):
    """ Scale aaomega spectra.

    spec              : List of Spectrum objects
    bscale and rscale : Arrays to multiply by the 385R and 580V
                        spectra, respectively.

    Returns
    -------
    spec  :  Input spectra with the 385R and 580V spectra scaled.
    grats :  List of gratings for each spectrum.
    """
    grats = [getgrat(s.wa) for s in spec]
    # multiply to match the J1201 templates
    sb = spec[grats.index('580V')]
    sr = spec[grats.index('385R')]
    sb.multiply(bscale)
    sr.multiply(rscale)
    # multiply red and blue halves to match in overlap region
    mult = astro.spec.scale_overlap(sb.wa, sb.fl, sb.er, sr.wa, sr.fl, sr.er)
    sr.multiply(mult)
    return spec, grats

def combine_scaled_580V_385R(specb, specr):
    """ Combine scaled spectra.

    Returns the combined spectrum.
    """
    # find the first good pixel in red spectrum
    i = 0
    while np.isnan(specr.fl[i]) or np.isnan(specr.er[i]):  i += 1
    sr = specr[i+2:]
    # make new wavelength scale
    ibmax = specb.wa.searchsorted(sr.wa[0])
    wa = np.concatenate( (specb.wa[:ibmax-1], sr.wa) )
    rsb = specb.rebin(wa=wa)
    rsr = sr.rebin(wa=wa)
    combined = astro.spec.combine( (rsb, rsr) )
    return combined

def scale_and_combine(qsos, spectra, bscale, rscale, pause=False):
    """ Scales and combines spectra, calling above two routines in the
    process."""
    fig = pl.gcf()
    wrapper = astro.wrappers.WrapPlotlines(fig)
    combined = []
    for i,(spec,info) in enumerate(spectra):
        if not spec: continue
        # scale the spectra to match one another.  First scale to
        # match template, then scale to match each other where they
        # overlap
        spec, grats = scaling(spec, bscale, rscale)
        i580V = grats.index('580V')
        i385R = grats.index('385R')
        ref = np.median(spec[i580V].fl[~np.isnan(spec[i580V].er)])
        for s,grat in zip(spec,grats):
            if grat in '1000R 1500Va 1500Vb'.split():
                s.multiply( abs(ref / np.median(s.fl[~np.isnan(s.er)])) )
        # plot spectra and the combination to check they look ok
        aaom.plotobject(spec, info, wmin=3800, wmax=8500, scale=0)
        a = pl.gca()
        wrapper.plotlines(qsos[i]['z'] + 1, a)
        a.set_title('%(name)s z=%(z)s mag=%(mag)s' % qsos[i])
        comb = combine_scaled_580V_385R(spec[i580V], spec[i385R])
        comb.plot(flcolor='k', ercolor='g',alpha=0.5)
        if pause:
            raw_input('enter to continue')
        combined.append( Bunch(sp=comb, fldinfo=(info[i580V], info[i385R]),
                               qsoinfo=qsos[i]) )
    wrapper.disconnect()
    return combined

def fit_cont_and_write_spec(dirname, combined, divmult=0.5):
    """ Write out combined spectra."""
    template = \
"""# name %(name)s
# ra    %(ra).10f    (j2000 deg)
# dec   %(dec).10f   (j2000 deg)
# z     %(z).3f
# mag   %(mag).3f    (B/V/R)
# fwhm  231          (instrumental broadening, km/s)
"""
    if not os.path.lexists(dirname):
        os.mkdir(dirname)
    for i,comb in enumerate(combined):
        print '%i of %i' % (i+1, len(combined)) 
        filename = os.path.join(dirname, '%s.txt' % comb.qsoinfo['name'])
        if os.path.lexists(filename):
            c = raw_input(
                '%s exists: Redo continuum and overwrite? ' % filename)
            if c.lower()[0] == 'n':  continue
        
        s = comb.sp
        z = comb.qsoinfo['z']
        if z < 0:  continue
        comb.sp.co,cpts = astro.fit.fitqsocont(s.wa, s.fl, s.er, z,
                                               divmult=divmult, debug=False)
        header = template % comb.qsoinfo
        note = raw_input('Note:').strip()
        header = header + '# note   %s\n' % note
        if note.startswith('bad'):
            comb.sp.co *= np.nan
        header = header + '# spline_nodes  %s\n' % repr(cpts)
        comb.sp.write(filename, header=header, overwrite=True)

    return combined

# for a given field (J0124, J1201, HE0940, Q0042)

# 1. read in all the separate r and b grating qso spectra.


# 2. find the scaling for each wavelength setting using a reference
# spectrum. Can be up to 5 wavelength settings (hires 2 b and 1 r,
# lores 1 b and 1 r).  Either by using a known object flux (from
# e.g. sloan) or bu using scaling relations from other fields.

# 3. Check the scaling applies to each qso. Adjust by hand if not.

# 4. Combine the scaled spectra to a single wavelength
# scale. (constant velocity scale? or rebin red scale to match blue
# scale?).  Probably want to keep lo and hir spectra separate.  Use
# lo/high res spectra to check wavelength calibration (also use sdss
# spectra to check wavelength calibration). Apply a shift if one is
# found.

# 5. Fit a continuum to the combined spectra.

# 6. Write out spectra

def process_he0940(fitcont_and_write=False):
    print 'Reading AAOmega spectra'
    PREFIX1 = '/home/nhmc/data/aaomega/feb08/combined/wigglez_comb/'
    PREFIX2 = '/home/nhmc/data/aaomega/march17/apr08_reduction/'
    b0 = aaom.readspec(PREFIX1 + 'he0940_580V.fits.gz')
    r0 = aaom.readspec(PREFIX1 + 'he0940_385R.fits.gz')
    b1 = aaom.readspec(PREFIX1 + 'he0940_1500V.fits.gz')
    r1 = aaom.readspec(PREFIX2 + 'he0940_1000R_6274.fits.gz')
    # j1201 scaling
    BSCALE = readtxt(PREFIX1 + '580V_scaling.txt', usecols=[1])
    RSCALE = readtxt(PREFIX1 + '385R_scaling.txt', usecols=[1])

    qsos = readqso_info('/home/nhmc/projects/gas_gal_hiz/qsos/qsos_0940.txt')
    spectra = match_spectra_to_qsos([b0,r0,b1,r1], qsos, arcsec=2)
    combined = scale_and_combine(qsos, spectra, BSCALE, RSCALE)

    if fitcont_and_write:
        combined = fit_cont_and_write_spec('he0940', combined, divmult=0.5)

    return combined

def process_j1201(fitcont_and_write=False):
    print 'Reading AAOmega spectra'
    PREFIX1 = '/home/nhmc/data/aaomega/feb08/combined/wigglez_comb/'
    PREFIX2 = '/home/nhmc/data/aaomega/march17/apr08_reduction/'
    b0 = aaom.readspec(PREFIX1 + 'j1201_580V.fits.gz')
    r0 = aaom.readspec(PREFIX1 + 'j1201_385R.fits.gz')
    b1 = aaom.readspec(PREFIX1 + 'j1201_1500V.fits.gz')
    # j1201 scaling
    BSCALE = readtxt(PREFIX1 + '580V_scaling.txt', usecols=[1])
    RSCALE = readtxt(PREFIX1 + '385R_scaling.txt', usecols=[1])

    qsos = readqso_info('/home/nhmc/projects/gas_gal_hiz/qsos/qsos_1201.txt')
    spectra = match_spectra_to_qsos([b0,r0,b1], qsos, arcsec=2)
    combined = scale_and_combine(qsos, spectra, BSCALE, RSCALE)

    if fitcont_and_write:
        combined = fit_cont_and_write_spec('j1201', combined, divmult=0.5)

    return combined

def process_pks2126(fitcont_and_write=False, pause=False):
    print 'Reading AAOmega spectra'
    prefix = '/home/nhmc/data/'
    b0 = aaom.readspec(prefix + 'aaomega/june08/pks2126_580V_4800.fits.gz')
    r0 = aaom.readspec(prefix + 'aaomega/june08/pks2126_385R_7300.fits.gz')
    b1a = aaom.readspec(prefix + 'aaomqso/pks2126_1500V_4625.fits.gz')
    b1b = aaom.readspec(prefix + 'aaomqso/pks2126_1500V_5365.fits.gz')
    r1 = aaom.readspec(prefix + 'aaomqso/pks2126_1000R_6280.fits.gz')

    qsos = readqso_info('/home/nhmc/projects/gas_gal_hiz/qsos/qsos_2126.txt')
    spectra = match_spectra_to_qsos([b0,r0,b1a,b1b,r1], qsos, arcsec=2)

    # use combination of worseck + j1201 scaling. Messy.
    BSCALE = readtxt(prefix + 'aaomega/feb08/combined/wigglez_comb/580V_scaling.txt', usecols=[1])
    RSCALE = readtxt(prefix + 'aaomega/feb08/combined/wigglez_comb/385R_scaling.txt', usecols=[1])

    calibspec = astro.spec.read('/home/nhmc/data/worseck_qsos/WW_2126/QSOJ21301-1533_cont_uc.dat')
    calibspec.multiply(1e16)
    grats = [getgrat(s.wa) for s in spectra[13][0]]
    i580V = grats.index('580V')
    i385R = grats.index('385R')
    bscale = astro.fit.poly_scale_spec(spectra[13][0][i580V],calibspec,order=7,debug=0)
    bscale2 = astro.fit.splice(BSCALE, bscale*25, 600, 1400)
    rscale = astro.fit.poly_scale_spec(spectra[13][0][i385R],calibspec,order=5,debug=0)
    rscale2 = astro.fit.splice(rscale*16.35, RSCALE, 250, 350)
    
    combined = scale_and_combine(qsos, spectra, bscale2, rscale2, pause=pause)

    if fitcont_and_write:
        combined = fit_cont_and_write_spec('pks2126', combined, divmult=0.7)

    return combined

def process_q0042(fitcont_and_write=False, pause=True):
    print 'Reading AAOmega spectra'
    prefix = '/home/nhmc/data/aaomqso/'
    r0 = aaom.readspec(PREFIX2 + 'q0042_1000R_6280.fits.gz')
    b0 = aaom.readspec(PREFIX2 + 'q0042_1500V_4625.fits.gz')
    b1 = aaom.readspec(PREFIX2 + 'q0042_1500V_5375.fits.gz')

    qsos = readqso_info('/home/nhmc/projects/gas_gal_hiz/qsos/qsos_0124.txt')
    spectra = match_spectra_to_qsos([b0,r0], qsos, arcsec=2)

    fig = pl.gcf()
    wrapper = astro.wrappers.WrapPlotlines(fig)
    combined = []
    for i,(spec,info) in enumerate(spectra):
        if not spec: continue
        # scale the spectra to match one another.  First scale to
        # match template, then scale to match each other where they
        # overlap
        grats = [getgrat(s.wa) for s in spec]
        # multiply to match the J1201 templates
        sb = spec[grats.index('580V')]
        sr = spec[grats.index('385R')]
        sb.multiply(bscale)
        sr.multiply(rscale)
        # multiply red and blue halves to match in overlap region
        mult = astro.spec.scale_overlap(sb.wa, sb.fl, sb.er, sr.wa, sr.fl, sr.er)
        sr.multiply(mult)
        
        i580V = grats.index('580V')
        i385R = grats.index('385R')
        ref = np.median(spec[i580V].fl[~np.isnan(spec[i580V].er)])
        for s,grat in zip(spec,grats):
            if grat in '1000R 1500Va 1500Vb'.split():
                s.multiply( abs(ref / np.median(s.fl[~np.isnan(s.er)])) )
        # plot spectra and the combination to check they look ok
        aaom.plotobject(spec, info, wmin=3800, wmax=8500, scale=0)
        a = pl.gca()
        wrapper.plotlines(qsos[i]['z'] + 1, a)
        a.set_title('%(name)s z=%(z)s mag=%(mag)s' % qsos[i])
        comb = combine_scaled_580V_385R(spec[i580V], spec[i385R])
        comb.plot(flcolor='k', ercolor='g',alpha=0.5)
        if pause:
            raw_input('enter to continue')
        combined.append( Bunch(sp=comb, fldinfo=(info[i580V], info[i385R]),
                               qsoinfo=qsos[i]) )
    wrapper.disconnect()
    if fitcont_and_write:
        combined = fit_cont_and_write_spec('q0042', combined, divmult=0.5)

    return combined



def process_j0124(fitcont_and_write=False, pause=True):
    # Trickier, because QSOs have different fibre numbers between observations
    print 'Reading AAOmega spectra'
    PREFIX1 = '/home/nhmc/data/aaomega/feb08/combined/wigglez_comb/'
    PREFIX2 = '/home/nhmc/data/aaomega/oct08/'
    #b0 = aaom.readspec(PREFIX2 + 'j0124_blue_081024.fits.gz')
    #r0 = aaom.readspec(PREFIX2 + 'j0124_red_081024.fits.gz')
    #b1 = aaom.readspec(PREFIX2 + 'j0124_combined_blue_081025.fits.gz')
    #r1 = aaom.readspec(PREFIX2 + 'j0124_combined_red_081025.fits.gz')
    cr = aaom.readspec(PREFIX2 + 'j0124_385R_7300.fits')
    cb = aaom.readspec(PREFIX2 + 'j0124_580V_4800.fits')

    qsos = readqso_info('/home/nhmc/projects/gas_gal_hiz/qsos/qsos_0124.txt')
    
    spectra = match_spectra_to_qsos([cr,cb], qsos, arcsec=2)
    #spectra = match_spectra_to_qsos([b0,r0,b1,r1], qsos, arcsec=2)
    # scaling
    refspec = astro.spec.read('/home/nhmc/data/sdss/j0124/1d_25/0399/1d/spSpec-51817-0399-375.fit')
    sp = spectra[40][0]
    rscale = astro.fit.poly_scale_spec(sp[0],refspec,debug=0,order=5, clip=5.7)
    bscale = astro.fit.poly_scale_spec(sp[1],refspec,debug=0,order=4)
    #BSCALE = readtxt(PREFIX1 + '580V_scaling.txt', usecols=[1])
    #RSCALE = readtxt(PREFIX1 + '385R_scaling.txt', usecols=[1])

    fig = pl.gcf()
    wrapper = astro.wrappers.WrapPlotlines(fig)
    combined = []
    for i,(spec,info) in enumerate(spectra):
        if not spec: continue
        # scale the spectra to match one another.  First scale to
        # match template, then scale to match each other where they
        # overlap
        grats = [getgrat(s.wa) for s in spec]
        # multiply to match the J1201 templates
        sb = spec[grats.index('580V')]
        sr = spec[grats.index('385R')]
        sb.multiply(bscale)
        sr.multiply(rscale)
        # multiply red and blue halves to match in overlap region
        mult = astro.spec.scale_overlap(sb.wa, sb.fl, sb.er, sr.wa, sr.fl, sr.er)
        sr.multiply(mult)
        i580V = grats.index('580V')
        i385R = grats.index('385R')
        # plot spectra and the combination to check they look ok
        aaom.plotobject(spec, info, wmin=3800, wmax=8500, scale=0)
        a = pl.gca()
        wrapper.plotlines(qsos[i]['z'] + 1, a)
        a.set_title('%(name)s z=%(z)s mag=%(mag)s' % qsos[i])
        comb = combine_scaled_580V_385R(spec[i580V], spec[i385R])
        comb.plot(flcolor='k', ercolor='g',alpha=0.5)
        pl.show()
        if pause:
            raw_input('enter to continue')
        combined.append( Bunch(sp=comb, fldinfo=(info[i580V], info[i385R]),
                               qsoinfo=qsos[i]) )
    wrapper.disconnect()

    if fitcont_and_write:
        combined = fit_cont_and_write_spec('j0124', combined, divmult=0.5)

    return combined


# if 0:    
#     # plot qso positions
#     pl.figure()
#     plot(qsos.ra,qsos.dec,'x')
#     plot([cra],[cdec],'s')
#     lbgs = gas_gal_hiz.read.lbgs('j0124')
#     plot(lbgs.ra,lbgs.dec,',')
#     foo = 25./60
#     x0,x1,y0,y1 = cra-foo, cra+foo, cdec-foo, cdec+foo
#     poly = pl.fill([x0,x0,x1,x1],[y0,y1,y1,y0],fc='None')
#     for q in cqsos:
#         pl.text(q['ra'],q['dec'],'%(mag)s %(z)s' % q)
#         pl.text(q['ra'],q['dec'],'%(mag)s %(z)s' % q)
#     pl.axis('equal')


#     b0,ib0 = aaom.readspec('j0124_blue_081024.fits.gz')
#     r0,ir0 = aaom.readspec('j0124_red_081024.fits.gz')
#     b1,ib1 = aaom.readspec('j0124_combined_blue_081025.fits.gz')
#     r1,ir1 = aaom.readspec('j0124_combined_red_081025.fits.gz')
#     filename = '/home/nhmc/data/aaomega/oct08/qsos_0124.txt'
